{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "05681fff-8d6a-48a3-a38f-ec6cab8b0418",
   "metadata": {},
   "source": [
    "# SplitRec：在隐语中使用拆分 BST 算法（Tensorflow 后端）"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4290e552-9930-40b1-bab8-4f24f3747608",
   "metadata": {},
   "source": [
    "阿里搜索推荐团队在2019年arXiv上发布文章《Behavior Sequence Transformer for E-commerce Recommendation in Alibaba》，提出BST模型。BST利用Transformer结构捕捉用户行为序列信息，解决WDL、DIN等模型忽略序列信息的问题，因而受到关注。本文将介绍如何在隐语中使用Tensorflow后端的拆分BST算法。"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e789ff9b-9f5e-466c-99a0-517b52572637",
   "metadata": {},
   "source": [
    "## BST 模型\n",
    "BST 模型在 DNN 模型基础上，使用 Transformer Layer 捕捉用户行为序列中的信息，整体结构如下\n",
    "![bst](./resources/bst0.png)\r\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "38d298bf-e5d1-45ce-b75b-143658ab18de",
   "metadata": {},
   "source": [
    "其中 Transformer Layer 使用一层 multi-head self-attention 结构，如图右侧所示。\n",
    "\n",
    "Transformer layer 的输入是序列特征的 Embedding 和位置 Embedding 的加和。"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a925168a-d700-4fd7-b3c4-73b4d8908a0a",
   "metadata": {},
   "source": [
    "## 隐语中的拆分 BST 模型"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c0242c4f-a72b-4a92-9736-38971146f27f",
   "metadata": {},
   "source": [
    "隐语考虑两方拆分学习场景，两方各持有一些特征，包括序列特征和其他特，一方持有 label 。双方特征经过 base 模型处理后，将 base 模型的输出输入到 fuse 模型进行学习，根据对序列特征的处理不同分为拆分 BST 模型和拆分 BSTPlus 模型。"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ba6f7a22-4e9a-4a50-aac4-b1071fb035ff",
   "metadata": {},
   "source": [
    "### 拆分 BST 模型\n",
    "其中序列特征通过 Transformer Layer 处理，其他特征通过普通的 Embedding 处理，序列特征和其他特征是可选的，模型结构如图所示。\n",
    "\n",
    "![bst](./resources/bst1.png)\n",
    "\n",
    "### 拆分 BSTPlus 模型\n",
    "上述 BST 模型虽然能够有效捕捉序列信息，但对于同一用户不同商品的预估来说，序列信息都是相同的，并不能区分同一用户行为序列与不同商品的关联性，因此 BSTPlus 模型在 BST 基础上叠加 DIN 中的 attention 思想，对 target item 和 Transformer layer 的输出进行 attention pooling，通过计算 target item 对序列中 item 的注意力来挖掘用户行为序列与 target item 之间的关联信息，结构如下图。由于 target item 存在于 label 方，暂不支持无 label 方有序列特征。\n",
    "\n",
    "![bst](./resources/bst2.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4096d2d3-8991-454a-8b67-53239860248e",
   "metadata": {},
   "source": [
    "## 隐语封装\n",
    "我们在隐语中提供了对于各种应用的封装。 关于 BST 的封装在 secretflow/ml/nn/applications/sl_bst_tf.py，提供了 `BSTBase`，`BSTPlusBase` 和 `BSTFuse` 几个类。\n",
    "下面我们通过一个例子来看一下如何使用隐语封装的 BST 来进行训练。"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a9955b23-eb62-4dee-b88d-be85ba50a154",
   "metadata": {},
   "source": [
    "## 环境设置"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "66f50a5c-563f-4fd8-a263-b480283da87d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The version of SecretFlow: 1.2.0.dev20231106\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2023-11-09 06:23:21,777\tINFO worker.py:1538 -- Started a local Ray instance.\n"
     ]
    }
   ],
   "source": [
    "import secretflow as sf\n",
    "\n",
    "# Check the version of your SecretFlow\n",
    "print('The version of SecretFlow: {}'.format(sf.__version__))\n",
    "\n",
    "# In case you have a running secretflow runtime already.\n",
    "sf.shutdown()\n",
    "sf.init(['alice', 'bob'], address=\"local\", log_to_driver=False)\n",
    "alice, bob = sf.PYU('alice'), sf.PYU('bob')\n",
    "device_y = bob"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ba03e5f6-b537-4bdc-adef-4b4f2ba7267b",
   "metadata": {},
   "source": [
    "## 数据集介绍\n",
    "这里将使用最经典的 MovieLens 数据集进行演示。 MovieLens 是一个开放式的推荐系统数据集，包含了电影评分和电影元数据信息。\n",
    "\n",
    "[数据集官网](https://grouplens.org/datasets/movielens/)\n",
    "\n",
    "[下载数据集](https://secretflow-data.oss-accelerate.aliyuncs.com/datasets/movielens/ml-1m.zip)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6726bbb9-a165-44fb-a6ac-6de652c0ff7b",
   "metadata": {},
   "source": [
    "## 下载并处理数据\n",
    "我们通过聚合同一用户按时间顺序评分过的电影生成用户行为序列，并对数据进行切分：  \n",
    "- alice: gender, age_group, occupation\n",
    "- bob: user_id, target_movie_id, sequence_movie_ids, label   "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "80a5cdb4-39c9-41f9-9613-7aaf5bb9b5bd",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from pathlib import Path\n",
    "import shutil\n",
    "\n",
    "data_dir = './test_bst_data_tf'\n",
    "dataset_download_dir = data_dir + '/data_download'\n",
    "gen_data_path = data_dir + '/data_sl_bst'\n",
    "\n",
    "\n",
    "def generate_data(plus=True):\n",
    "    import numpy as np\n",
    "    import pandas as pd\n",
    "\n",
    "    from secretflow.utils.simulation.datasets import _DATASETS, get_dataset, unzip\n",
    "\n",
    "    global dataset_download_dir\n",
    "    if not Path(dataset_download_dir).is_dir():\n",
    "        filepath = get_dataset(_DATASETS['ml-1m'])\n",
    "        unzip(filepath, dataset_download_dir)\n",
    "\n",
    "    dataset_dir = dataset_download_dir + '/ml-1m'\n",
    "    # read data\n",
    "    users = pd.read_csv(\n",
    "        dataset_dir + \"/users.dat\",\n",
    "        sep=\"::\",\n",
    "        names=[\"user_id\", \"gender\", \"age_group\", \"occupation\", \"zip_code\"],\n",
    "        engine='python',\n",
    "    )\n",
    "\n",
    "    ratings = pd.read_csv(\n",
    "        dataset_dir + \"/ratings.dat\",\n",
    "        sep=\"::\",\n",
    "        names=[\"user_id\", \"movie_id\", \"rating\", \"unix_timestamp\"],\n",
    "        engine='python',\n",
    "    )\n",
    "\n",
    "    movies = pd.read_csv(\n",
    "        dataset_dir + \"/movies.dat\",\n",
    "        sep=\"::\",\n",
    "        names=[\"movie_id\", \"title\", \"genres\"],\n",
    "        engine='python',\n",
    "        encoding=\"ISO-8859-1\",\n",
    "    )\n",
    "\n",
    "    users[\"user_id\"] = users[\"user_id\"].apply(lambda x: f\"{x}\")\n",
    "    users[\"age_group\"] = users[\"age_group\"].apply(lambda x: f\"{x}\")\n",
    "    users[\"occupation\"] = users[\"occupation\"].apply(lambda x: f\"{x}\")\n",
    "\n",
    "    movies[\"movie_id\"] = movies[\"movie_id\"].apply(lambda x: f\"{x}\")\n",
    "    movies[\"genres\"] = movies[\"genres\"].apply(lambda x: ','.join(x.split('|')))\n",
    "\n",
    "    ratings[\"movie_id\"] = ratings[\"movie_id\"].apply(lambda x: f\"{x}\")\n",
    "    ratings[\"user_id\"] = ratings[\"user_id\"].apply(lambda x: f\"{x}\")\n",
    "    ratings[\"rating\"] = ratings[\"rating\"].apply(lambda x: float(x))\n",
    "\n",
    "    # generate users' movie sequence\n",
    "    ratings_group = ratings.sort_values(by=[\"unix_timestamp\"]).groupby(\"user_id\")\n",
    "\n",
    "    ratings_data = pd.DataFrame(\n",
    "        data={\n",
    "            \"user_id\": list(ratings_group.groups.keys()),\n",
    "            \"movie_ids\": list(ratings_group.movie_id.apply(list)),\n",
    "            \"ratings\": list(ratings_group.rating.apply(list)),\n",
    "            \"timestamps\": list(ratings_group.unix_timestamp.apply(list)),\n",
    "        }\n",
    "    )\n",
    "\n",
    "    # split sequence\n",
    "    if plus:\n",
    "        sequence_length = 5\n",
    "        step_size = 4\n",
    "    else:\n",
    "        sequence_length = 4\n",
    "        step_size = 2\n",
    "\n",
    "    def create_sequences(values, window_size, step_size):\n",
    "        sequences = []\n",
    "        start_index = 0\n",
    "        while True:\n",
    "            end_index = start_index + window_size\n",
    "            seq = values[start_index:end_index]\n",
    "            if len(seq) < window_size:\n",
    "                seq.extend(['[PAD]'] * (window_size - len(seq)))\n",
    "                sequences.append(seq)\n",
    "                break\n",
    "            sequences.append(seq)\n",
    "            start_index += step_size\n",
    "        return sequences\n",
    "\n",
    "    ratings_data.movie_ids = ratings_data.movie_ids.apply(\n",
    "        lambda ids: create_sequences(ids, sequence_length, step_size)\n",
    "    )\n",
    "\n",
    "    ratings_data.ratings = ratings_data.ratings.apply(\n",
    "        lambda ids: create_sequences(ids, sequence_length, step_size)\n",
    "    )\n",
    "    del ratings_data[\"timestamps\"]\n",
    "\n",
    "    ratings_data_movies = ratings_data[[\"user_id\", \"movie_ids\"]].explode(\n",
    "        \"movie_ids\", ignore_index=True\n",
    "    )\n",
    "    ratings_data_rating = ratings_data[[\"ratings\"]].explode(\n",
    "        \"ratings\", ignore_index=True\n",
    "    )\n",
    "    ratings_data_transformed = pd.concat(\n",
    "        [ratings_data_movies, ratings_data_rating], axis=1\n",
    "    )\n",
    "    ratings_data_transformed = ratings_data_transformed.join(\n",
    "        users.set_index(\"user_id\"), on=\"user_id\"\n",
    "    )\n",
    "    ratings_data_transformed['movie_id'] = ratings_data_transformed.movie_ids.apply(\n",
    "        lambda x: x[-1] if '[PAD]' not in x else x[x.index('[PAD]') - 1]\n",
    "    )\n",
    "\n",
    "    if plus:\n",
    "        # last movie in sequence is target movie\n",
    "        ratings_data_transformed.movie_ids = ratings_data_transformed.movie_ids.apply(\n",
    "            lambda x: (\n",
    "                \",\".join(x[:-1])\n",
    "                if '[PAD]' not in x\n",
    "                else \",\".join(x[: x.index('[PAD]') - 1] + x[x.index('[PAD]') :])\n",
    "            )\n",
    "        )\n",
    "    else:\n",
    "        ratings_data_transformed.movie_ids = ratings_data_transformed.movie_ids.apply(\n",
    "            lambda x: \",\".join(x)\n",
    "        )\n",
    "\n",
    "    ratings_data_transformed['label'] = ratings_data_transformed.ratings.apply(\n",
    "        lambda x: x[-1] if '[PAD]' not in x else x[x.index('[PAD]') - 1]\n",
    "    )\n",
    "\n",
    "    ratings_data_transformed.ratings = ratings_data_transformed.ratings.apply(\n",
    "        lambda x: \",\".join([str(v) for v in x[:-1]])\n",
    "    )\n",
    "    ratings_data_transformed = ratings_data_transformed.join(\n",
    "        movies.set_index(\"movie_id\"), on=\"movie_id\"\n",
    "    )\n",
    "\n",
    "    del (\n",
    "        ratings_data_transformed[\"zip_code\"],\n",
    "        ratings_data_transformed[\"title\"],\n",
    "        ratings_data_transformed[\"genres\"],\n",
    "        ratings_data_transformed[\"ratings\"],\n",
    "    )\n",
    "\n",
    "    ratings_data_transformed.rename(\n",
    "        columns={\"movie_ids\": \"sequence_movie_ids\", \"movie_id\": \"target_movie_id\"},\n",
    "        inplace=True,\n",
    "    )\n",
    "    random_selection = np.random.rand(len(ratings_data_transformed.index)) <= 0.85\n",
    "    train_data = ratings_data_transformed[random_selection]\n",
    "    test_data = ratings_data_transformed[~random_selection]\n",
    "\n",
    "    # save data\n",
    "    if os.path.exists(gen_data_path):\n",
    "        shutil.rmtree(gen_data_path)\n",
    "    os.mkdir(gen_data_path)\n",
    "    os.mkdir(gen_data_path + '/vocabulary')\n",
    "\n",
    "    train_data.to_csv(\n",
    "        gen_data_path + \"/train_data.csv\", index=False, sep=\"|\", encoding='utf-8'\n",
    "    )\n",
    "    test_data.to_csv(\n",
    "        gen_data_path + \"/test_data.csv\", index=False, sep=\"|\", encoding='utf-8'\n",
    "    )\n",
    "\n",
    "    train_data_alice = train_data[[\"gender\", \"age_group\", \"occupation\"]]\n",
    "    train_data_bob = train_data[\n",
    "        [\n",
    "            \"user_id\",\n",
    "            \"sequence_movie_ids\",\n",
    "            \"target_movie_id\",\n",
    "            \"label\",\n",
    "        ]\n",
    "    ]\n",
    "\n",
    "    test_data_alice = test_data[[\"gender\", \"age_group\", \"occupation\"]]\n",
    "    test_data_bob = test_data[\n",
    "        [\n",
    "            \"user_id\",\n",
    "            \"sequence_movie_ids\",\n",
    "            \"target_movie_id\",\n",
    "            \"label\",\n",
    "        ]\n",
    "    ]\n",
    "\n",
    "    train_data_alice.to_csv(\n",
    "        gen_data_path + \"/train_data_alice.csv\", index=False, sep=\"|\", encoding='utf-8'\n",
    "    )\n",
    "    train_data_bob.to_csv(\n",
    "        gen_data_path + \"/train_data_bob.csv\", index=False, sep=\"|\", encoding='utf-8'\n",
    "    )\n",
    "\n",
    "    test_data_alice.to_csv(\n",
    "        gen_data_path + \"/test_data_alice.csv\", index=False, sep=\"|\", encoding='utf-8'\n",
    "    )\n",
    "    test_data_bob.to_csv(\n",
    "        gen_data_path + \"/test_data_bob.csv\", index=False, sep=\"|\", encoding='utf-8'\n",
    "    )\n",
    "\n",
    "    with open(gen_data_path + '/vocabulary/user_id', 'w') as f:\n",
    "        f.write('\\n'.join(list(ratings_data_transformed.user_id.unique())))\n",
    "\n",
    "    with open(gen_data_path + '/vocabulary/gender', 'w') as f:\n",
    "        f.write('\\n'.join(list(ratings_data_transformed.gender.unique())))\n",
    "\n",
    "    with open(gen_data_path + '/vocabulary/age_group', 'w') as f:\n",
    "        f.write('\\n'.join(list(ratings_data_transformed.age_group.unique())))\n",
    "\n",
    "    with open(gen_data_path + '/vocabulary/occupation', 'w') as f:\n",
    "        f.write('\\n'.join(list(ratings_data_transformed.occupation.unique())))\n",
    "\n",
    "    with open(gen_data_path + '/vocabulary/item_id', 'w') as f:\n",
    "        f.write('\\n'.join(list(movies.movie_id.unique())))\n",
    "\n",
    "\n",
    "generate_data()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "6652f32f-3099-4c1a-9b58-dfdedd27e10c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "gender|age_group|occupation\n",
      "F|1|10\n",
      "F|1|10\n",
      "F|1|10\n",
      "F|1|10\n",
      "F|1|10\n",
      "F|1|10\n",
      "F|1|10\n",
      "F|1|10\n",
      "F|1|10\n"
     ]
    }
   ],
   "source": [
    "! head ./test_bst_data_tf/data_sl_bst/train_data_alice.csv"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "a6598709-0ddb-46cc-8cb2-77891470281d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "user_id|sequence_movie_ids|target_movie_id|label\n",
      "1|3186,1721,1270,1022|2340|3.0\n",
      "1|2340,1836,3408,1207|2804|5.0\n",
      "1|2804,260,720,1193|919|4.0\n",
      "1|2028,3105,938,1035|1962|4.0\n",
      "1|1962,1028,2018,150|1097|4.0\n",
      "1|1097,914,1287,2797|1246|4.0\n",
      "1|1246,2762,661,2918|531|4.0\n",
      "1|531,3114,2791,1029|2321|3.0\n",
      "1|2321,1197,594,2398|1545|4.0\n"
     ]
    }
   ],
   "source": [
    "! head ./test_bst_data_tf/data_sl_bst/train_data_bob.csv"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b144729b-e17a-4676-af63-32274f97de17",
   "metadata": {},
   "source": [
    "到这里，我们以及产出训练测试用的数据文件。"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e436048e-0abd-431d-b3b4-4e80b4511a11",
   "metadata": {},
   "source": [
    "## 构造 dataset_builder 处理数据\n",
    "这里对特征进行 ID 化处理，rating > 3 作为正样本，否则作为负样本 "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "9f9419fa-02f3-45a9-b1be-330b877e8fce",
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_dataset_builder_alice(\n",
    "    batch_size=128,\n",
    "    repeat_count=5,\n",
    "):\n",
    "    def dataset_builder(x):\n",
    "        import pandas as pd\n",
    "        import tensorflow as tf\n",
    "\n",
    "        x = [dict(t) if isinstance(t, pd.DataFrame) else t for t in x]\n",
    "        x = x[0] if len(x) == 1 else tuple(x)\n",
    "        data_set = (\n",
    "            tf.data.Dataset.from_tensor_slices(x).batch(batch_size).repeat(repeat_count)\n",
    "        )\n",
    "\n",
    "        return data_set\n",
    "\n",
    "    return dataset_builder\n",
    "\n",
    "\n",
    "def create_dataset_builder_bob(\n",
    "    batch_size=128,\n",
    "    repeat_count=5,\n",
    "):\n",
    "    def _parse_bob(row_sample, label):\n",
    "        import tensorflow as tf\n",
    "\n",
    "        y_t = label[\"label\"]\n",
    "        y = tf.expand_dims(\n",
    "            tf.where(\n",
    "                y_t > 3,\n",
    "                tf.ones_like(y_t, dtype=tf.float32),\n",
    "                tf.zeros_like(y_t, dtype=tf.float32),\n",
    "            ),\n",
    "            axis=1,\n",
    "        )\n",
    "        return row_sample, y\n",
    "\n",
    "    def dataset_builder(x):\n",
    "        import pandas as pd\n",
    "        import tensorflow as tf\n",
    "\n",
    "        x = [dict(t) if isinstance(t, pd.DataFrame) else t for t in x]\n",
    "        x = x[0] if len(x) == 1 else tuple(x)\n",
    "        data_set = (\n",
    "            tf.data.Dataset.from_tensor_slices(x).batch(batch_size).repeat(repeat_count)\n",
    "        )\n",
    "\n",
    "        data_set = data_set.map(_parse_bob)\n",
    "\n",
    "        return data_set\n",
    "\n",
    "    return dataset_builder\n",
    "\n",
    "\n",
    "batch_size = 128\n",
    "epoch = 1\n",
    "data_builder_dict = {\n",
    "    alice: create_dataset_builder_alice(\n",
    "        batch_size=batch_size,\n",
    "        repeat_count=epoch,\n",
    "    ),\n",
    "    bob: create_dataset_builder_bob(\n",
    "        batch_size=batch_size,\n",
    "        repeat_count=epoch,\n",
    "    ),\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5a731615-01d8-4227-b6b6-f2b7b2894bbc",
   "metadata": {},
   "source": [
    "## 定义模型结构\n",
    "通过 plus 指定 base 模型结构\n",
    "- plus=True: BSTBase\n",
    "- plus=False: BSTPlusBase"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "95a67d44-443a-4e44-b8b2-36775de21094",
   "metadata": {},
   "outputs": [],
   "source": [
    "from secretflow.ml.nn.applications.sl_bst_tf import BSTBase, BSTFuse, BSTPlusBase\n",
    "\n",
    "\n",
    "def create_bstplus_base_model_alice():\n",
    "    # Create model\n",
    "    def create_model():\n",
    "        import tensorflow as tf\n",
    "        from tensorflow.keras import layers\n",
    "\n",
    "        fea_list = [\"gender\", \"age_group\", \"occupation\"]\n",
    "        fea_emb_size = {\"gender\": 8, \"age_group\": 8, \"occupation\": 8}\n",
    "        fea_voc = {}\n",
    "        for key in fea_list:\n",
    "            with open(gen_data_path + '/vocabulary/' + key) as f:\n",
    "                values = [line.strip() for line in f.readlines()]\n",
    "                fea_voc[key] = values\n",
    "\n",
    "        def preprocess():\n",
    "            inputs = {}\n",
    "            outputs = {}\n",
    "            for key in fea_list:\n",
    "                inputs[key] = tf.keras.Input(shape=(1,), dtype=tf.string)\n",
    "\n",
    "                idx = layers.StringLookup(\n",
    "                    vocabulary=fea_voc[key],\n",
    "                    mask_token=None,\n",
    "                    num_oov_indices=0,\n",
    "                    name=f\"{key}_index_lookup\",\n",
    "                )(inputs[key])\n",
    "\n",
    "                outputs[key] = layers.Embedding(\n",
    "                    input_dim=len(fea_voc[key]),\n",
    "                    output_dim=fea_emb_size[key],\n",
    "                    name=f\"{key}_embedding\",\n",
    "                )(idx)\n",
    "\n",
    "            return tf.keras.Model(inputs=inputs, outputs=outputs)\n",
    "\n",
    "        preprocess_layer = preprocess()\n",
    "        model = BSTPlusBase(\n",
    "            preprocess_layer=preprocess_layer,\n",
    "            dnn_units_size=[32],\n",
    "        )\n",
    "        model.compile(\n",
    "            loss=tf.keras.losses.binary_crossentropy,\n",
    "            optimizer=tf.keras.optimizers.Adam(),\n",
    "            metrics=[\n",
    "                tf.keras.metrics.AUC(),\n",
    "                tf.keras.metrics.Precision(),\n",
    "                tf.keras.metrics.Recall(),\n",
    "            ],\n",
    "        )\n",
    "        return model  # need wrap\n",
    "\n",
    "    return create_model\n",
    "\n",
    "\n",
    "def create_bstplus_base_model_bob():\n",
    "    # Create model\n",
    "    def create_model():\n",
    "        import tensorflow as tf\n",
    "        from tensorflow.keras import layers\n",
    "\n",
    "        fea_list = [\"user_id\"]\n",
    "        fea_emb_size = {\"user_id\": 8}\n",
    "        fea_voc = {}\n",
    "        for key in fea_list:\n",
    "            with open(gen_data_path + '/vocabulary/' + key) as f:\n",
    "                values = [line.strip() for line in f.readlines()]\n",
    "                fea_voc[key] = values\n",
    "\n",
    "        with open(gen_data_path + '/vocabulary/item_id') as f:\n",
    "            values = [line.strip() for line in f.readlines()]\n",
    "            fea_voc[\"item\"] = values\n",
    "\n",
    "        seq_len = 4\n",
    "        item_embedding_dims = 9\n",
    "\n",
    "        def preprocess():\n",
    "            inputs = {}\n",
    "            outputs = {}\n",
    "            for key in fea_list:\n",
    "                inputs[key] = tf.keras.Input(shape=(1,), dtype=tf.string)\n",
    "\n",
    "                idx = layers.StringLookup(\n",
    "                    vocabulary=fea_voc[key],\n",
    "                    mask_token=None,\n",
    "                    num_oov_indices=0,\n",
    "                    name=f\"{key}_index_lookup\",\n",
    "                )(inputs[key])\n",
    "\n",
    "                outputs[key] = layers.Embedding(\n",
    "                    input_dim=len(fea_voc[key]),\n",
    "                    output_dim=fea_emb_size[key],\n",
    "                    name=f\"{key}_embedding\",\n",
    "                )(idx)\n",
    "\n",
    "            # sequence input\n",
    "            inputs[\"sequence_movie_ids\"] = tf.keras.Input(\n",
    "                name=\"sequence_movie_ids\",\n",
    "                shape=(1,),\n",
    "                dtype=tf.string,\n",
    "            )\n",
    "\n",
    "            inputs['target_movie_id'] = tf.keras.Input(\n",
    "                name=\"target_movie_id\",\n",
    "                shape=(1,),\n",
    "                dtype=tf.string,\n",
    "            )\n",
    "\n",
    "            item_lookup_layer = layers.StringLookup(\n",
    "                vocabulary=fea_voc[\"item\"],\n",
    "                mask_token='[PAD]',  # note here!\n",
    "                name=\"item_index_lookup\",\n",
    "            )\n",
    "\n",
    "            seq_split = tf.strings.split(inputs[\"sequence_movie_ids\"], ',').to_tensor(\n",
    "                '[PAD]', shape=[None, 1, 4]\n",
    "            )\n",
    "\n",
    "            item_idx = item_lookup_layer(seq_split)\n",
    "            outputs[\"sequence_idx\"] = item_idx\n",
    "\n",
    "            target_item_idx = item_lookup_layer(inputs['target_movie_id'])\n",
    "            outputs['target_movie_id'] = target_item_idx\n",
    "\n",
    "            return tf.keras.Model(inputs=inputs, outputs=outputs)\n",
    "\n",
    "        preprocess_layer = preprocess()\n",
    "        model = BSTPlusBase(\n",
    "            preprocess_layer=preprocess_layer,\n",
    "            dnn_units_size=[32],\n",
    "            sequence_fea=[\"sequence_idx\"],\n",
    "            target_fea=\"target_movie_id\",\n",
    "            item_embedding_dims={\"target_movie_id\": item_embedding_dims},\n",
    "            seq_len={\"sequence_idx\": seq_len},\n",
    "            item_voc_size={\"target_movie_id\": len(fea_voc[\"item\"])},\n",
    "            num_head={\"sequence_idx\": 3},\n",
    "            dropout_rate={\"sequence_idx\": 0.1},\n",
    "        )\n",
    "\n",
    "        model.compile(\n",
    "            loss=tf.keras.losses.binary_crossentropy,\n",
    "            optimizer=tf.keras.optimizers.Adam(),\n",
    "            metrics=[\n",
    "                tf.keras.metrics.AUC(),\n",
    "                tf.keras.metrics.Precision(),\n",
    "                tf.keras.metrics.Recall(),\n",
    "            ],\n",
    "        )\n",
    "        return model  # need wrap\n",
    "\n",
    "    return create_model\n",
    "\n",
    "\n",
    "def create_base_model_alice():\n",
    "    # Create model\n",
    "    def create_model():\n",
    "        import tensorflow as tf\n",
    "        from tensorflow.keras import layers\n",
    "\n",
    "        fea_list = [\"gender\", \"age_group\", \"occupation\"]\n",
    "        fea_emb_size = {\"gender\": 8, \"age_group\": 8, \"occupation\": 8}\n",
    "        fea_voc = {}\n",
    "        for key in fea_list:\n",
    "            with open(gen_data_path + '/vocabulary/' + key) as f:\n",
    "                values = [line.strip() for line in f.readlines()]\n",
    "                fea_voc[key] = values\n",
    "\n",
    "        def preprocess():\n",
    "            inputs = {}\n",
    "            outputs = {}\n",
    "            for key in fea_list:\n",
    "                inputs[key] = tf.keras.Input(shape=(1,), dtype=tf.string)\n",
    "\n",
    "                idx = layers.StringLookup(\n",
    "                    vocabulary=fea_voc[key],\n",
    "                    mask_token=None,\n",
    "                    num_oov_indices=0,\n",
    "                    name=f\"{key}_index_lookup\",\n",
    "                )(inputs[key])\n",
    "\n",
    "                outputs[key] = layers.Embedding(\n",
    "                    input_dim=len(fea_voc[key]),\n",
    "                    output_dim=fea_emb_size[key],\n",
    "                    name=f\"{key}_embedding\",\n",
    "                )(idx)\n",
    "\n",
    "            return tf.keras.Model(inputs=inputs, outputs=outputs)\n",
    "\n",
    "        preprocess_layer = preprocess()\n",
    "        model = BSTBase(\n",
    "            preprocess_layer=preprocess_layer,\n",
    "            dnn_units_size=[32],\n",
    "        )\n",
    "        model.compile(\n",
    "            loss=tf.keras.losses.binary_crossentropy,\n",
    "            optimizer=tf.keras.optimizers.Adam(),\n",
    "            metrics=[\n",
    "                tf.keras.metrics.AUC(),\n",
    "                tf.keras.metrics.Precision(),\n",
    "                tf.keras.metrics.Recall(),\n",
    "            ],\n",
    "        )\n",
    "        return model  # need wrap\n",
    "\n",
    "    return create_model\n",
    "\n",
    "\n",
    "def create_base_model_bob():\n",
    "    # Create model\n",
    "    def create_model():\n",
    "        import tensorflow as tf\n",
    "        from tensorflow.keras import layers\n",
    "\n",
    "        fea_list = [\"user_id\"]\n",
    "        fea_emb_size = {\"user_id\": 8}\n",
    "        fea_voc = {}\n",
    "        for key in fea_list:\n",
    "            with open(gen_data_path + '/vocabulary/' + key) as f:\n",
    "                values = [line.strip() for line in f.readlines()]\n",
    "                fea_voc[key] = values\n",
    "\n",
    "        with open(gen_data_path + '/vocabulary/item_id') as f:\n",
    "            values = [line.strip() for line in f.readlines()]\n",
    "            fea_voc[\"item\"] = values\n",
    "\n",
    "        seq_len = 4\n",
    "        item_embedding_dims = 9\n",
    "\n",
    "        def preprocess():\n",
    "            inputs = {}\n",
    "            outputs = {}\n",
    "            for key in fea_list:\n",
    "                inputs[key] = tf.keras.Input(shape=(1,), dtype=tf.string)\n",
    "\n",
    "                idx = layers.StringLookup(\n",
    "                    vocabulary=fea_voc[key],\n",
    "                    mask_token=None,\n",
    "                    num_oov_indices=0,\n",
    "                    name=f\"{key}_index_lookup\",\n",
    "                )(inputs[key])\n",
    "\n",
    "                outputs[key] = layers.Embedding(\n",
    "                    input_dim=len(fea_voc[key]),\n",
    "                    output_dim=fea_emb_size[key],\n",
    "                    name=f\"{key}_embedding\",\n",
    "                )(idx)\n",
    "\n",
    "            # sequence input\n",
    "            inputs[\"sequence_movie_ids\"] = tf.keras.Input(\n",
    "                name=\"sequence_movie_ids\",\n",
    "                shape=(1,),\n",
    "                dtype=tf.string,\n",
    "            )\n",
    "\n",
    "            seq_split = tf.strings.split(inputs[\"sequence_movie_ids\"], ',').to_tensor(\n",
    "                '[PAD]', shape=[None, 1, 4]\n",
    "            )\n",
    "\n",
    "            item_idx = layers.StringLookup(\n",
    "                vocabulary=fea_voc[\"item\"],\n",
    "                mask_token='[PAD]',  # note here!\n",
    "                name=\"item_index_lookup\",\n",
    "            )(seq_split)\n",
    "            outputs[\"sequence_idx\"] = item_idx\n",
    "\n",
    "            return tf.keras.Model(inputs=inputs, outputs=outputs)\n",
    "\n",
    "        preprocess_layer = preprocess()\n",
    "        model = BSTBase(\n",
    "            preprocess_layer=preprocess_layer,\n",
    "            dnn_units_size=[32],\n",
    "            sequence_fea=[\"sequence_idx\"],\n",
    "            item_embedding_dims={\"sequence_idx\": item_embedding_dims},\n",
    "            seq_len={\"sequence_idx\": seq_len},\n",
    "            item_voc_size={\"sequence_idx\": len(fea_voc[\"item\"])},\n",
    "            num_head={\"sequence_idx\": 2},\n",
    "            dropout_rate={\"sequence_idx\": 0.5},\n",
    "        )\n",
    "\n",
    "        model.compile(\n",
    "            loss=tf.keras.losses.binary_crossentropy,\n",
    "            optimizer=tf.keras.optimizers.Adam(),\n",
    "            metrics=[\n",
    "                tf.keras.metrics.AUC(),\n",
    "                tf.keras.metrics.Precision(),\n",
    "                tf.keras.metrics.Recall(),\n",
    "            ],\n",
    "        )\n",
    "        return model  # need wrap\n",
    "\n",
    "    return create_model\n",
    "\n",
    "\n",
    "def create_fuse_model():\n",
    "    # Create model\n",
    "    def create_model():\n",
    "        import tensorflow as tf\n",
    "\n",
    "        model = BSTFuse(dnn_units_size=[256, 128])\n",
    "        model.compile(\n",
    "            loss=tf.keras.losses.binary_crossentropy,\n",
    "            optimizer=tf.keras.optimizers.Adam(),\n",
    "            metrics=[\n",
    "                tf.keras.metrics.AUC(),\n",
    "                tf.keras.metrics.Precision(),\n",
    "                tf.keras.metrics.Recall(),\n",
    "            ],\n",
    "        )\n",
    "        return model\n",
    "\n",
    "    return create_model\n",
    "\n",
    "\n",
    "plus = True\n",
    "if plus:\n",
    "    model_base_alice = create_bstplus_base_model_alice()\n",
    "    model_base_bob = create_bstplus_base_model_bob()\n",
    "else:\n",
    "    model_base_alice = create_base_model_alice()\n",
    "    model_base_bob = create_base_model_bob()\n",
    "\n",
    "base_model_dict = {\n",
    "    alice: model_base_alice,\n",
    "    bob: model_base_bob,\n",
    "}\n",
    "model_fuse = create_fuse_model()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f0f6c821-896e-4507-b9bb-edbb7e47e3a3",
   "metadata": {},
   "source": [
    "## 定义 SL Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "662950ad-18b4-4fb6-bd41-625d347d16e8",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:root:Create proxy actor <class 'secretflow.ml.nn.sl.backend.tensorflow.sl_base.PYUSLTFModel'> with party alice.\n",
      "INFO:root:Create proxy actor <class 'secretflow.ml.nn.sl.backend.tensorflow.sl_base.PYUSLTFModel'> with party bob.\n"
     ]
    }
   ],
   "source": [
    "from secretflow.ml.nn import SLModel\n",
    "\n",
    "sl_model = SLModel(\n",
    "    base_model_dict=base_model_dict,\n",
    "    device_y=device_y,\n",
    "    model_fuse=model_fuse,\n",
    "    random_seed=0,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6e8e085e-cc66-45c4-8afa-04331d4bba86",
   "metadata": {},
   "source": [
    "## 开始训练"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "46ea8414-e0bf-4e1d-9ba6-0c125c540538",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:root:Create proxy actor <class 'secretflow.data.core.agent.PartitionAgent'> with party alice.\n",
      "INFO:root:Create proxy actor <class 'secretflow.data.core.agent.PartitionAgent'> with party bob.\n",
      "INFO:root:SL Train Params: {'x': VDataFrame(partitions={PYURuntime(alice): <secretflow.data.core.partition.Partition object at 0x7faafcce22b0>, PYURuntime(bob): <secretflow.data.core.partition.Partition object at 0x7faafccdf9d0>}, aligned=True), 'y': VDataFrame(partitions={PYURuntime(bob): <secretflow.data.core.partition.Partition object at 0x7faafccdfac0>}, aligned=True), 'batch_size': 128, 'epochs': 1, 'verbose': 1, 'callbacks': None, 'validation_data': (VDataFrame(partitions={PYURuntime(alice): <secretflow.data.core.partition.Partition object at 0x7faafcce22b0>, PYURuntime(bob): <secretflow.data.core.partition.Partition object at 0x7faafccdf9d0>}, aligned=True), VDataFrame(partitions={PYURuntime(bob): <secretflow.data.core.partition.Partition object at 0x7faafccdfac0>}, aligned=True)), 'shuffle': False, 'sample_weight': None, 'validation_freq': 1, 'dp_spent_step_freq': None, 'dataset_builder': {PYURuntime(alice): <function create_dataset_builder_alice.<locals>.dataset_builder at 0x7faafd07eaf0>, PYURuntime(bob): <function create_dataset_builder_bob.<locals>.dataset_builder at 0x7faafd07ec10>}, 'audit_log_params': {}, 'random_seed': 0, 'audit_log_dir': None, 'self': <secretflow.ml.nn.sl.sl_model.SLModel object at 0x7fabc4fa5fa0>}\n",
      "100%|██████████| 1675/1675 [09:27<00:00,  2.95it/s, epoch: 1/1 -  train_loss:0.6462613940238953  train_auc_1:0.6504741907119751  train_precision_1:0.6380547881126404  train_recall_1:0.8145814538002014  val_loss:0.6162347197532654  val_auc_1:0.7021257281303406  val_precision_1:0.6727387309074402  val_recall_1:0.8101728558540344 ]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'train_loss': [0.6462614], 'train_auc_1': [0.6504742], 'train_precision_1': [0.6380548], 'train_recall_1': [0.81458145], 'val_loss': [0.6162347], 'val_auc_1': [0.7021257], 'val_precision_1': [0.67273873], 'val_recall_1': [0.81017286]}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "from secretflow.data.vertical import read_csv\n",
    "\n",
    "\n",
    "vdf = read_csv(\n",
    "    {\n",
    "        alice: gen_data_path + '/train_data_alice.csv',\n",
    "        bob: gen_data_path + '/train_data_bob.csv',\n",
    "    },\n",
    "    delimiter='|',\n",
    ")\n",
    "label = vdf[\"label\"]\n",
    "data = vdf.drop(columns=[\"label\"])\n",
    "\n",
    "data[\"user_id\"] = data[\"user_id\"].astype(\"string\")\n",
    "data[\"age_group\"] = data[\"age_group\"].astype(\"string\")\n",
    "data[\"occupation\"] = data[\"occupation\"].astype(\"string\")\n",
    "data[\"target_movie_id\"] = data[\"target_movie_id\"].astype(\"string\")\n",
    "\n",
    "history = sl_model.fit(\n",
    "    data,\n",
    "    label,\n",
    "    validation_data=(data, label),\n",
    "    epochs=epoch,\n",
    "    batch_size=batch_size,\n",
    "    random_seed=0,\n",
    "    dataset_builder=data_builder_dict,\n",
    ")\n",
    "print(history)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fc6fc68c-6071-4a8b-b2e8-0005bb4ab68d",
   "metadata": {},
   "source": [
    "## 总结\n",
    "本文通过 MovieLens 数据集上的推荐任务来演示了如何通过隐语来训练拆分 BST 模型，您需要\n",
    "1. 下载并处理数据集；\n",
    "2. 构造 dataset_builder 处理数据；\r",
    "3. 定义模型结构，调用 `BSTBase`，`BSTPlusBase`和`BSTFuse`定义行模结构；\n",
    "4. 使用 SLModel 进行训练，预测，评估即可。"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3635c590-fdc6-4283-810e-469c28e4d3f7",
   "metadata": {},
   "source": [
    "您可以在自己的数据集上进行尝试，如有任何问题，可以在 github 上进行讨论。"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.17"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
